#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <gmssl/sm4.h>
#include <gmssl/rand.h>
#include <gmssl/error.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>
#include "../../include/crypto/sm4_crypto.h"

#define BUFFER_SIZE 4096

// 安全清除敏感数据
static void secure_zero(void *ptr, size_t size)
{
    volatile unsigned char *p = ptr;
    while (size--)
        *p++ = 0;
}

// 生成随机SM4密钥
int generate_sm4_key(unsigned char *key)
{
    if (!key)
        return -1;

    return rand_bytes(key, SM4_KEY_SIZE) == 1 ? 0 : -1;
}

// 生成随机初始化向量
int generate_sm4_iv(unsigned char *iv)
{
    if (!iv)
        return -1;

    return rand_bytes(iv, SM4_IV_SIZE) == 1 ? 0 : -1;
}

// SM4-CBC模式加密
int sm4_cbc_encrypt_data(const unsigned char *plaintext, size_t plaintext_len,
                         const unsigned char *key, const unsigned char *iv,
                         unsigned char *ciphertext, size_t *ciphertext_len)
{
    if (!plaintext || !key || !iv || !ciphertext || !ciphertext_len)
    {
        return -1;
    }

    SM4_CBC_CTX cbc_ctx;
    int ret = -1;

    // 估算密文长度（考虑PKCS7填充）
    // 密文长度 = ((明文长度 / 块大小) + 1) * 块大小
    size_t padded_len = ((plaintext_len + SM4_BLOCK_SIZE - 1) / SM4_BLOCK_SIZE + 1) * SM4_BLOCK_SIZE;

    if (*ciphertext_len < padded_len)
    {
        *ciphertext_len = padded_len;
        return -2; // 输出缓冲区太小
    }

    // 初始化CBC模式加密 - 直接使用密钥
    if (sm4_cbc_encrypt_init(&cbc_ctx, key, iv) != 1)
    {
        return -1;
    }

    // 处理数据
    size_t outlen = 0;
    if (sm4_cbc_encrypt_update(&cbc_ctx, plaintext, plaintext_len,
                               ciphertext, &outlen) != 1)
    {
        goto cleanup;
    }

    // 处理最后填充
    size_t finish_outlen = 0;
    if (sm4_cbc_encrypt_finish(&cbc_ctx, ciphertext + outlen, &finish_outlen) != 1)
    {
        goto cleanup;
    }

    *ciphertext_len = outlen + finish_outlen;
    ret = 0;

cleanup:
    // 清除敏感数据
    secure_zero(&cbc_ctx, sizeof(SM4_CBC_CTX));
    return ret;
}

// SM4-CBC模式解密
int sm4_cbc_decrypt_data(const unsigned char *ciphertext, size_t ciphertext_len,
                         const unsigned char *key, const unsigned char *iv,
                         unsigned char *plaintext, size_t *plaintext_len)
{
    if (!ciphertext || !key || !iv || !plaintext || !plaintext_len)
    {
        return -1;
    }

    // 密文长度必须是块大小的整数倍
    if (ciphertext_len % SM4_BLOCK_SIZE != 0)
    {
        fprintf(stderr, "SM4-CBC解密失败: 密文长度不是块大小的整数倍\n");
        return -1;
    }

    if (*plaintext_len < ciphertext_len)
    {
        *plaintext_len = ciphertext_len;
        return -2; // 输出缓冲区太小
    }

    SM4_CBC_CTX cbc_ctx;
    int ret = -1;

    // 初始化CBC模式解密 - 直接使用密钥
    if (sm4_cbc_decrypt_init(&cbc_ctx, key, iv) != 1)
    {
        return -1;
    }

    // 解密数据
    size_t outlen = 0;
    if (sm4_cbc_decrypt_update(&cbc_ctx, ciphertext, ciphertext_len,
                               plaintext, &outlen) != 1)
    {
        goto cleanup;
    }

    // 处理最后填充
    size_t finish_outlen = 0;
    if (sm4_cbc_decrypt_finish(&cbc_ctx, plaintext + outlen, &finish_outlen) != 1)
    {
        goto cleanup;
    }

    *plaintext_len = outlen + finish_outlen;
    ret = 0;

cleanup:
    // 清除敏感数据
    secure_zero(&cbc_ctx, sizeof(SM4_CBC_CTX));
    return ret;
}

// SM4-CBC模式加密文件
int sm4_encrypt_file(const char *input_file, const char *output_file,
                     const unsigned char *key, const unsigned char *iv)
{
    if (!input_file || !output_file || !key || !iv)
    {
        return -1;
    }

    int input_fd = open(input_file, O_RDONLY);
    if (input_fd == -1)
    {
        perror("打开输入文件失败");
        return -1;
    }

    int output_fd = open(output_file, O_WRONLY | O_CREAT | O_TRUNC, 0644);
    if (output_fd == -1)
    {
        perror("创建输出文件失败");
        close(input_fd);
        return -1;
    }

    // 写入IV到文件开头
    if (write(output_fd, iv, SM4_IV_SIZE) != SM4_IV_SIZE)
    {
        perror("写入IV失败");
        close(input_fd);
        close(output_fd);
        return -1;
    }

    // 加密文件内容
    SM4_CBC_CTX cbc_ctx;
    unsigned char buffer_in[BUFFER_SIZE];
    unsigned char buffer_out[BUFFER_SIZE + SM4_BLOCK_SIZE]; // 额外空间用于填充
    int bytes_read;
    int ret = -1;

    // 初始化CBC模式加密
    if (sm4_cbc_encrypt_init(&cbc_ctx, key, iv) != 1)
    {
        close(input_fd);
        close(output_fd);
        return -1;
    }

    // 读取并加密数据块
    while ((bytes_read = read(input_fd, buffer_in, BUFFER_SIZE)) > 0)
    {
        size_t outlen = sizeof(buffer_out);

        // 处理完整块
        if (sm4_cbc_encrypt_update(&cbc_ctx, buffer_in, bytes_read,
                                   buffer_out, &outlen) != 1)
        {
            goto cleanup_file;
        }

        // 写入加密的完整块
        if (write(output_fd, buffer_out, outlen) != outlen)
        {
            perror("写入加密数据失败");
            goto cleanup_file;
        }
    }

    if (bytes_read == -1)
    {
        perror("读取输入文件失败");
        goto cleanup_file;
    }

    // 处理最后的填充
    size_t finish_outlen = sizeof(buffer_out);
    if (sm4_cbc_encrypt_finish(&cbc_ctx, buffer_out, &finish_outlen) != 1)
    {
        goto cleanup_file;
    }

    // 写入最后填充块
    if (finish_outlen > 0 && write(output_fd, buffer_out, finish_outlen) != finish_outlen)
    {
        perror("写入填充数据失败");
        goto cleanup_file;
    }

    ret = 0; // 成功

cleanup_file:
    // 清除敏感数据
    secure_zero(&cbc_ctx, sizeof(SM4_CBC_CTX));
    secure_zero(buffer_in, sizeof(buffer_in));
    secure_zero(buffer_out, sizeof(buffer_out));

    close(input_fd);
    close(output_fd);

    return ret;
}

// SM4-CBC模式解密文件
int sm4_decrypt_file(const char *input_file, const char *output_file,
                     const unsigned char *key, const unsigned char *iv)
{
    if (!input_file || !output_file || !key || !iv)
    {
        return -1;
    }

    int input_fd = open(input_file, O_RDONLY);
    if (input_fd == -1)
    {
        perror("打开输入文件失败");
        return -1;
    }

    // 获取文件大小
    struct stat st;
    if (fstat(input_fd, &st) == -1)
    {
        perror("获取文件信息失败");
        close(input_fd);
        return -1;
    }

    // 检查文件大小是否合法
    if (st.st_size < SM4_IV_SIZE || (st.st_size - SM4_IV_SIZE) % SM4_BLOCK_SIZE != 0)
    {
        fprintf(stderr, "无效的加密文件: 大小不正确\n");
        close(input_fd);
        return -1;
    }

    int output_fd = open(output_file, O_WRONLY | O_CREAT | O_TRUNC, 0644);
    if (output_fd == -1)
    {
        perror("创建输出文件失败");
        close(input_fd);
        return -1;
    }

    // 读取文件中的IV
    unsigned char file_iv[SM4_IV_SIZE];
    if (read(input_fd, file_iv, SM4_IV_SIZE) != SM4_IV_SIZE)
    {
        perror("读取IV失败");
        close(input_fd);
        close(output_fd);
        return -1;
    }

    // 解密文件内容
    SM4_CBC_CTX cbc_ctx;
    unsigned char buffer_in[BUFFER_SIZE + SM4_BLOCK_SIZE]; // 确保块对齐
    unsigned char buffer_out[BUFFER_SIZE + SM4_BLOCK_SIZE];
    int bytes_read;
    int ret = -1;

    // 初始化CBC模式解密，使用从文件读取的IV
    if (sm4_cbc_decrypt_init(&cbc_ctx, key, iv ? iv : file_iv) != 1)
    {
        close(input_fd);
        close(output_fd);
        return -1;
    }

    // 读取并解密数据块
    while ((bytes_read = read(input_fd, buffer_in, BUFFER_SIZE)) > 0)
    {
        size_t outlen = sizeof(buffer_out);

        // 解密数据块
        if (sm4_cbc_decrypt_update(&cbc_ctx, buffer_in, bytes_read,
                                   buffer_out, &outlen) != 1)
        {
            goto cleanup_file;
        }

        // 写入解密数据
        if (outlen > 0 && write(output_fd, buffer_out, outlen) != outlen)
        {
            perror("写入解密数据失败");
            goto cleanup_file;
        }
    }

    if (bytes_read == -1)
    {
        perror("读取输入文件失败");
        goto cleanup_file;
    }

    // 处理最后的填充
    size_t finish_outlen = sizeof(buffer_out);
    if (sm4_cbc_decrypt_finish(&cbc_ctx, buffer_out, &finish_outlen) != 1)
    {
        goto cleanup_file;
    }

    // 写入最后处理的数据（移除填充）
    if (finish_outlen > 0 && write(output_fd, buffer_out, finish_outlen) != finish_outlen)
    {
        perror("写入解密数据失败");
        goto cleanup_file;
    }

    ret = 0; // 成功

cleanup_file:
    // 清除敏感数据
    secure_zero(&cbc_ctx, sizeof(SM4_CBC_CTX));
    secure_zero(buffer_in, sizeof(buffer_in));
    secure_zero(buffer_out, sizeof(buffer_out));

    close(input_fd);
    close(output_fd);

    return ret;
}

// 安全清除密钥
void secure_clear_key(unsigned char *key, size_t key_len)
{
    if (key && key_len > 0)
    {
        secure_zero(key, key_len);
    }
}